import os
from PIL import Image  # 使用PIL库来处理图片

# 设置输入和输出路径
input_file = 'data/starai_datasets/data/starai_datasets/train_GT.txt'
output_dir = 'task/label'
image_dir = 'task/train'  # 图片文件存放的路径

# 如果输出目录不存在，则创建
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 读取原始标注文件
with open(input_file, 'r') as f:
    lines = f.readlines()

# 使用字典存储每帧的标注信息
annotations = {}

# 处理每一行
for line in lines:
    video_name, frame, bb_left, bb_top, bb_width, bb_height, class_id = line.strip().split(',')
    frame = int(frame)
    class_id = int(class_id) - 1
    
    # 构建图片文件路径
    image_name = f"{os.path.splitext(video_name)[0]}_frame_{frame:04d}.jpg"
    image_path = os.path.join(image_dir, image_name)

    # 尝试读取图片的分辨率
    try:
        with Image.open(image_path) as img:
            video_width, video_height = img.size
    except FileNotFoundError:
        print(f"未找到图片文件: {image_path}，跳过该帧。")
        continue  # 跳过该条记录

    # 计算中心点坐标和归一化
    bb_x_center = (int(bb_left) + int(bb_width) / 2) / video_width
    bb_y_center = (int(bb_top) + int(bb_height) / 2) / video_height
    bb_width_normalized = int(bb_width) / video_width
    bb_height_normalized = int(bb_height) / video_height

    # 创建帧的标注信息
    key = (video_name, frame)
    if key not in annotations:
        annotations[key] = []
    
    annotations[key].append(f"{class_id} {bb_x_center} {bb_y_center} {bb_width_normalized} {bb_height_normalized}")

# 写入标注文件
for (video_name, frame), ann in annotations.items():
    output_file = os.path.join(output_dir, f"{os.path.splitext(video_name)[0]}_frame_{frame:04d}.txt")
    
    # 仅在有标注时写入文件
    with open(output_file, 'a') as out_f:
        for line in ann:
            out_f.write(line + '\n')

print("标注文件生成完成。")